import pyaudio
import wave
import datetime
import numpy as np
import matplotlib.pyplot as plt
from collections import deque

# 音频输入参数
format = pyaudio.paInt16
channels = 1
rate = 44100
chunk_size = 1024
record_seconds = 2
pool_size = 50  # 池子的大小

# 创建一个PyAudio对象
audio = pyaudio.PyAudio()

# 打开音频输入流
stream = audio.open(format=format,
                    channels=channels,
                    rate=rate,
                    input=True,
                    frames_per_buffer=chunk_size)

# 初始化池子
sample_pool = deque(maxlen=pool_size)

def calculate_threshold(audio_data):
    if len(audio_data) == 0:
        return 0

    pool_average = np.nanmean(audio_data)  # 使用 np.nanmean() 来计算平均值，忽略 NaN 值
    if np.isnan(pool_average):
        return 0

    hit_threshold = int(0.5 * pool_average)
    return hit_threshold

def count_hits(audio_data, threshold):
    hit_count = 0
    skip_count = 0

    for sample in audio_data:
        if skip_count > 0:
            skip_count -= 1
            continue
        
        sample_pool.append(sample)
        pool_average = np.mean(sample_pool)

        if sample > threshold:
            print(f"HIT: {sample}")
            hit_count += 1
            skip_count = pool_size
            
        # 计算新的阈值
        threshold = int(0.5 * pool_average)

    return hit_count

def save_audio(filename, frames):
    with wave.open(filename, "wb") as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(audio.get_sample_size(format))
        wf.setframerate(rate)
        wf.writeframes(b''.join(frames))

def plot_waveform(audio_data):
    # 创建时间轴
    time_axis = np.arange(0, len(audio_data)) / rate

    # 绘制波形图
    plt.figure()
    plt.plot(time_axis, audio_data)
    plt.xlabel("Time (s)")
    plt.ylabel("Amplitude")
    plt.title("Waveform")
    plt.show()

input("按回车键开始录制音频...")

print("录音开始...")
while True:
    frames = []
    for i in range(int(rate / chunk_size * record_seconds)):
        data = stream.read(chunk_size)
        frames.append(data)

    current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"audio_{current_time}.wav"
    save_audio(filename, frames)
    print(f"音频已保存为：{filename}")

    with wave.open(filename, "rb") as wf:
        num_frames = wf.getnframes()
        audio_data = np.frombuffer(wf.readframes(num_frames), dtype=np.int16)
    
    hit_threshold = calculate_threshold(audio_data)
    hit_count = count_hits(audio_data, hit_threshold)
    print(f"HIT次数：{hit_count}")
    
    plot_waveform(audio_data)

    user_input = input("按回车键继续录制音频，输入其他字符退出程序：")
    if user_input != "":
        break

    print("录音开始...")

print("录音结束...")
